package edu.utah.seq.analysis;

import java.io.*;
import java.util.regex.*;
import java.util.*;

import util.bio.annotation.*;
import util.bio.parsers.*;
import util.gen.*;
import edu.utah.seq.data.*;
import edu.utah.seq.parsers.*;
import trans.tpmap.*;

/**Takes PointData and generates sliding window scan statistics. Unlike the original ScanSeqs, this looks for a skew between strands.
 * @author nix
 * */
public class DefinedRegionScanSeqs {

	//user defined fields
	private File[] treatmentPointDirs;
	private File[] controlPointDirs;
	private File saveDirectory;
	private File fullPathToR = new File ("/usr/bin/R");
	private int peakShift = 0;
	private int halfPeakShift = 0;
	private File bedFile;
	private File refSeqFile;

	//internal fields
	private double numberTreatmentObservations;
	private double numberControlObservations;
	private double expectedFractionUp;
	private double expectedFractionDown;
	private double scalarTC;
	private double scalarCT;
	private HashMap<String, PointData[]> treatments;
	private HashMap<String, PointData[]> controls;
	private int numberCachedBinPVals = 500;
	private float[][] pValsDataUp;
	private float[][] pValsDataDown;
	private String chromosome;
	private HashMap<String,UCSCGeneLine[]> geneModels;
	private UCSCGeneLine[] geneLines;

	//constructors
	public DefinedRegionScanSeqs(String[] args){
		long startTime = System.currentTimeMillis();

		//set fields
		processArgs(args);

		//load gene models from refFlat for refSeq UCSC gene table
		UCSCGeneModelTableReader reader = null;
		if (refSeqFile != null){
			reader = new UCSCGeneModelTableReader(refSeqFile, 0);
			reader.splitByChromosome();
			geneModels = reader.getChromSpecificGeneLines();
			geneLines = reader.getGeneLines();
		}
		//or from bed file
		else if (bedFile != null) {
			Bed[] bed = Bed.parseFile(bedFile, 0, 0);
			geneLines = new UCSCGeneLine[bed.length];
			boolean addName = bed[0].getName().trim().equals("");
			for (int i=0; i< bed.length; i++){
				if (addName) bed[i].setName((i+1)+"");
				geneLines[i] = new UCSCGeneLine(bed[i]);
			}
			reader = new UCSCGeneModelTableReader();
			reader.setGeneLines(geneLines);
			reader.splitByChromosome();
			geneModels = reader.getChromSpecificGeneLines();
		}
		if (geneModels == null | geneModels.size() == 0) Misc.printExit("\nProblem loading your USCS gene model table or bed file.\n");
		//check ordering
		if (reader.checkStartStopOrder() == false) Misc.printExit("\nOne of your regions's coordinates are reversed. Check that each start is less than the stop.\n");
		//fetch data but don't load it.
		System.out.println("Fetching data...");
		fetchData();

		//make binomial pvalue matrix
		if (controls != null){
			System.out.println("Building binomial p-value look up tables...");
			expectedFractionUp = numberTreatmentObservations/(numberTreatmentObservations+numberControlObservations);
			expectedFractionDown = 1- expectedFractionUp;		
			pValsDataUp = Num.convertToFloat(Num.binomialPValMatrix(numberCachedBinPVals, expectedFractionUp, saveDirectory, fullPathToR, true));
			pValsDataDown = Num.convertToFloat(Num.binomialPValMatrix(numberCachedBinPVals, expectedFractionDown, saveDirectory, fullPathToR, true));
			scalarTC = numberTreatmentObservations/ numberControlObservations;
			scalarCT = numberControlObservations/ numberTreatmentObservations;	
			System.out.println("\tExpectUp "+expectedFractionUp+"\n\tExpectDown "+expectedFractionDown+"\n\tScalarTC "+scalarTC+"\n\tScalarCT "+scalarCT);
		}

		//for each chromosome in gene line data
		System.out.print("Scanning Annotated Chromosomes");
		Iterator<String> it = geneModels.keySet().iterator();
		int index = 0;
		while (it.hasNext()){
			//fetch chromosome
			chromosome = it.next();
			//fetch, and shift positions from the treatment and control
			if (halfPeakShift !=0) shiftPositions();

			//calculate binomial p-values 
			if (controls !=null) {
				calculateBinomialPValues();
			}
			else sumTreatmentReads();
			System.out.print(".");
		}
		System.out.println();

		if (controls != null){
			//convert binomial pvalues to FDRs
			System.out.println("Converting binomial p-values to q-value FDRs...");
			convertPValuesToQValues();
			//calc log2Ratios
			calculateLog2Ratios();
		}

		//sort
		Arrays.sort(geneLines, new UCSCGeneLineComparatorScore(0));

		//print
		printGeneModels();

		//finish and calc run time
		double diffTime = ((double)(System.currentTimeMillis() -startTime))/1000;
		System.out.println("\nDone! "+Math.round(diffTime)+" seconds\n");

	}

	public void printGeneModels(){
		try {
			File results = new File (saveDirectory, "drssResults.xls");
			PrintWriter out = new PrintWriter (new FileWriter( results));
			//general info
			String genome;
			PointData[] pd = treatments.values().iterator().next();
			if (pd[0]!=null) genome = pd[0].getInfo().getVersionedGenome();
			else genome = pd[1].getInfo().getVersionedGenome();
			
			String misc = "GenomeVersion="+genome+ ", TotalTreatObs="+(int)numberTreatmentObservations;
			String url = "=HYPERLINK(\"http://localhost:7085/UnibrowControl?version="+genome+"&seqid=";
			if (geneLines[0].getDisplayName() != null) out.print("#DisplayName\tName\tChr\tStrand\tStart\tStop\t");
			else out.print("#Name\tChr\tStrand\tStart\tStop\t");

			if (controls != null){
				misc = misc+ ", TotalCtrlObs="+(int)numberControlObservations;
				out.println("pVal\tqValFDR\tlog2((tSum+1)/(cSum+1))\ttSum+\ttSum-\tcSum+\tcSum-\t"+misc);
				for (int i=0; i< geneLines.length; i++){
					String name;
					if (geneLines[i].getDisplayName() !=null) name = geneLines[i].getDisplayName();
					else name = geneLines[i].getName();
					//url
					int start = geneLines[i].getTxStart() - 5000;
					if (start < 0) start = 0;
					int end = geneLines[i].getTxEnd() + 5000;
					out.print(url+geneLines[i].getChrom()+"&start="+start+"&end="+end+"\",\""+name+"\")\t");
					//print second name?
					if (geneLines[i].getDisplayName() !=null) out.print(geneLines[i].getName()+"\t");

					//scores = pVal,tSumPlus, tSumMinus, cSumPlus, cSumMinus,qVal,log2(ratio)
					float[] s = geneLines[i].getScores();
					out.println(geneLines[i].coordinates()+"\t"+s[0]+"\t"+ s[5]+"\t"+ s[6]+"\t"+ s[1]+"\t"+ s[2]+"\t"+ s[3]+"\t"+ s[4]);
				}
			}

			else{
				out.println("Sum\tSum+\tSum-\t"+misc);
				for (int i=0; i< geneLines.length; i++){
					String name;
					if (geneLines[i].getDisplayName() !=null) name = geneLines[i].getDisplayName();
					else name = geneLines[i].getName();
					//url
					int start = geneLines[i].getTxStart() - 5000;
					if (start < 0) start = 0;
					int end = geneLines[i].getTxEnd() + 5000;
					out.print(url+geneLines[i].getChrom()+"&start="+start+"&end="+end+"\",\""+name+"\")\t");
					//print second name?
					if (geneLines[i].getDisplayName() !=null) out.print(geneLines[i].getName()+"\t");
					float[] s = geneLines[i].getScores();
					out.println(geneLines[i].coordinates()+"\t"+s[0]+"\t"+ s[1]+"\t"+ s[2]);
				}
			}
			out.close();
		} catch (Exception e){
			e.printStackTrace();
		}
	}

	public void calculateLog2Ratios(){
		for (int i=0; i< geneLines.length; i++){
			float[] scores = geneLines[i].getScores();
			double tSum = scores[1] + scores[2];
			double cSum = scores[3] + scores[4];
			if (tSum == 0 && cSum == 0) scores[6] = 0;
			else scores[6] = calculateLog2Ratio(tSum, cSum);
		}
	}

	/**Calculates a log2( (tSum+1)/(cSum+1) ) on linearly scaled tSum and cSum based on the total observations.*/
	public float calculateLog2Ratio( double tSum, double cSum){
		double t;
		double c;
		if (tSum !=0 ) {
			t = tSum * scalarCT;
			c = cSum;
		}
		else {
			c = cSum * scalarTC;
			t = tSum;
		}
		double ratio = (t+1)/(c+1);
		return (float)Num.log2(ratio);
	}


	/**Converts binomial pvalues in SmoothingWindowInfo to qvalues using Storey's method.
	 * Only converts those with > 9 reads. */
	private void convertPValuesToQValues(){
		//collect those for qval
		ArrayList<UCSCGeneLine> geneLinesForQVal = new ArrayList<UCSCGeneLine>();
		for (int i=0; i< geneLines.length; i++){
			//scores = pVal,tSumPlus, tSumMinus, cSumPlus, cSumMinus,pValUp,pValDown
			float[] s = geneLines[i].getScores();
			float sum = s[1]+ s[2]+ s[3]+ s[4];
			if (sum > 9) geneLinesForQVal.add(geneLines[i]);
			else {
				s[5] = 0;
				s[6] = 0;
			}
		}
		int num = geneLinesForQVal.size();
		//collect pvalues 
		float[] up = new float[num];
		float[] down = new float[num];
		for (int i=0; i< num; i++){
			//scores = pVal,tSumPlus, tSumMinus, cSumPlus, cSumMinus,pValUp,pValDown
			float[] scores = geneLinesForQVal.get(i).getScores();
			up[i] = scores[5];
			down[i] = scores[6];
		}
		//calc qvals
		float[] upQVals = Num.qValueFDR(saveDirectory, up, fullPathToR, true, false);
		float[] downQVals = Num.qValueFDR(saveDirectory, down, fullPathToR, true, false);

		if (upQVals == null || downQVals == null){
			System.out.println("\tSetting qValue FDRs to zero.");
			//zero scores
			for (int i=0; i< num; i++){
				//scores = pVal,tSumPlus, tSumMinus, cSumPlus, cSumMinus,qVal,logRatio
				float[] scores = geneLinesForQVal.get(i).getScores();
				scores[5] = 0;
			}
		}
		else {
			//assign to regions
			for (int i=0; i< num; i++){
				
				//scores = pVal,tSumPlus, tSumMinus, cSumPlus, cSumMinus,qVal,logRatio
				float[] scores = geneLinesForQVal.get(i).getScores();
				if (upQVals[i] > downQVals[i]) scores[5] = upQVals[i];
				else scores[5] = -1 * downQVals[i];
			}
		}
		upQVals = null;
		downQVals = null;
	}

	private void fetchData(){
		System.out.println("Loading PointData...");
		treatments = PointData.fetchStrandedCombinePointData(treatmentPointDirs);
		numberTreatmentObservations = PointData.totalObservationsMultiPointData(treatments);
		System.out.println("\t"+(int)numberTreatmentObservations+" Treatment Observations");
		if (controlPointDirs != null){
			controls = PointData.fetchStrandedCombinePointData(controlPointDirs);
			numberControlObservations = PointData.totalObservationsMultiPointData(controls);
			System.out.println("\t"+(int)numberControlObservations+" Control Observations");
		}
	}


	/**For cases when just interested in treatment data.*/
	private void sumTreatmentReads(){
		//fetch PointData
		PointData[] pd = treatments.get(chromosome);
		PointData treatPlusPD = pd[0];
		PointData treatMinusPD = pd[1];

		//get UCSCGeneLine[]
		UCSCGeneLine[] geneLines = geneModels.get(chromosome);

		//for each gene line 
		for (int i=0; i< geneLines.length; i++){
			//calculate totals
			ExonIntron[] exons = geneLines[i].getExons();
			//fetch scores
			float tSumPlus = 0; 
			float tSumMinus = 0;
			for (int j=0; j< exons.length; j++){
				int start = exons[j].getStart();
				int stop = exons[j].getEnd();
				if (treatPlusPD != null) tSumPlus += treatPlusPD.sumPositionBP(start, stop); 
				if (treatMinusPD != null)tSumMinus += treatMinusPD.sumPositionBP(start, stop);
			}
			//scores = tSum, tSumPlus, tSumMinus
			float[] scores = new float[]{tSumPlus+tSumMinus,tSumPlus, tSumMinus};
			//save scores
			geneLines[i].setScores(scores);
		}
	}


	private void calculateBinomialPValues(){
		//fetch PointData
		PointData[] pd = treatments.get(chromosome);
		PointData treatPlusPD = pd[0];
		PointData treatMinusPD = pd[1];
		pd = controls.get(chromosome);
		PointData ctrlPlusPD = pd[0];
		PointData ctrlMinusPD = pd[1];

		//make arrays to calculate binomial p-vals with R
		ArrayList<UCSCGeneLine> forBinom = new ArrayList<UCSCGeneLine>();

		//get UCSCGeneLine[]
		UCSCGeneLine[] geneLines = geneModels.get(chromosome);

		//for each window 
		for (int i=0; i< geneLines.length; i++){
			//calculate totals
			ExonIntron[] exons = geneLines[i].getExons();
			//fetch scores
			float tSumPlus = 0; 
			float tSumMinus = 0;
			float cSumPlus = 0;
			float cSumMinus = 0;
			for (int j=0; j< exons.length; j++){
				int start = exons[j].getStart();
				int stop = exons[j].getEnd();
				if (treatPlusPD != null) tSumPlus += treatPlusPD.sumPositionBP(start, stop); 
				if (treatMinusPD != null)tSumMinus += treatMinusPD.sumPositionBP(start, stop);
				if (ctrlPlusPD != null)cSumPlus += ctrlPlusPD.sumPositionBP(start, stop);
				if (ctrlMinusPD != null)cSumMinus += ctrlMinusPD.sumPositionBP(start, stop);
			}

			//scores = pVal,tSumPlus, tSumMinus, cSumPlus, cSumMinus,qVal,log2(ratio)
			float[] scores = new float[]{0,tSumPlus, tSumMinus, cSumPlus, cSumMinus, 0, 0};

			//save scores
			geneLines[i].setScores(scores);
			float sum = tSumPlus+ tSumMinus+ cSumPlus+ cSumMinus;
			if (sum == 0) continue;

			//calc diff binomial from cache? or directly?
			float tSum = tSumPlus+ tSumMinus;
			float cSum = cSumPlus+ cSumMinus;

			//yes
			if (tSum < numberCachedBinPVals && cSum < numberCachedBinPVals){
				//calc pvals
				int t = (int)tSum;
				int c = (int)cSum;
				float up = pValsDataUp[t][c];
				float down = pValsDataDown[c][t];
				if (up > down) scores[0] = up;
				else scores[0] = -1* down;
				scores[5] = up;
				scores[6] = down;
			}
			//no
			else forBinom.add(geneLines[i]);
		}
		//calc binom from R
		if (forBinom.size()!=0) calculateBinomialsFromR(forBinom);
	}

	/**Takes UCSCGeneLine that had observations that exceeded cached pvalues and calculates
	 * directly from R.*/
	private void calculateBinomialsFromR(ArrayList<UCSCGeneLine> sm){
		int num = sm.size();
		//for each win, want to calculate up, down, skew
		int[][] upObs = new int[num][2];
		int[][] downObs = new int[num][2];
		for (int i=0; i< num; i++){
			//scores = tSumPlus, tSumMinus, cSumPlus, cSumMinus,pVal,qVal,log2(ratio)
			float[] scores = sm.get(i).getScores();
			int tSumPlus = (int)scores[1];
			int tSumMinus = (int)scores[2];
			int tSum = tSumPlus + tSumMinus;
			int cSumPlus = (int)scores[3];
			int cSumMinus = (int)scores[4];
			int cSum = cSumPlus + cSumMinus;
			//set obs
			upObs[i] = new int[]{tSum,cSum};
			downObs[i] = new int[]{cSum,tSum};
		}
		//fetch pvalues
		double[] upPvals = Num.binomialPValues(expectedFractionUp, saveDirectory, upObs, fullPathToR, true);
		double[] downPvals = Num.binomialPValues(expectedFractionDown, saveDirectory, downObs, fullPathToR, true);
		if (upPvals == null || downPvals == null || upPvals.length != num || downPvals.length != num) {
			Misc.printErrAndExit("\nProblem fetching binomial pvalues from R. Check that R really is at "+fullPathToR);
		}
		//assign scores
		for (int i=0; i< num; i++){
			float[] scores = sm.get(i).getScores();
			//scores = tSumPlus, tSumMinus, cSumPlus, cSumMinus,pVal,qVal,log2(ratio)
			scores[5] = (float)upPvals[i];
			scores[6] = (float)downPvals[i];
			//assign pVal score
			if (scores[5]> scores[6]) scores[0] = scores[5];
			else scores[0] = -1 * scores[6];
		}
	}

	/**Shifts the positions halfPeakShift (+ for sense, - for antisense) sets the positions into the data*/
	private void shiftPositions(){
		//fetch data from treatments
		PointData[] t = treatments.get(chromosome);
		if (t[0]!=null) {
			int[] p = t[0].getPositions();
			addShift(p,halfPeakShift);
		}
		if (t[1]!=null) {
			int[] p = t[1].getPositions();
			addShift(p, -1*halfPeakShift);
		}
		//fetch data from controls
		if (controls != null) {
			PointData[] c = controls.get(chromosome);
			if (c[0]!=null) {
				int[] p = c[0].getPositions();
				addShift(p, halfPeakShift);
			}
			if (c[1]!=null){
				int[] p = c[1].getPositions();
				addShift(p, -1*halfPeakShift);
			}
		}
	}

	/**Adds the toAdd to each int.*/
	public static void addShift(int[] positions, int toAdd){
		for (int i=0; i< positions.length; i++){
			positions[i] += toAdd;
			if (positions[i]<0) positions[i] = 0;
		}
	}

	public static void main(String[] args) {
		if (args.length ==0){
			printDocs();
			System.exit(0);
		}
		new DefinedRegionScanSeqs(args);
	}		


	/**This method will process each argument and assign new variables*/
	public void processArgs(String[] args){
		Pattern pat = Pattern.compile("-[a-z]");
		System.out.println("\nArguments: "+Misc.stringArrayToString(args, " ")+"\n");
		for (int i = 0; i<args.length; i++){
			String lcArg = args[i].toLowerCase();
			Matcher mat = pat.matcher(lcArg);
			if (mat.matches()){
				char test = args[i].charAt(1);
				try{
					switch (test){
					case 't': treatmentPointDirs = IO.extractFiles(args[++i]); break;
					case 'c': controlPointDirs = IO.extractFiles(args[++i]); break;
					case 's': saveDirectory = new File(args[++i]); break;
					case 'r': fullPathToR = new File(args[++i]); break;
					case 'u': refSeqFile = new File(args[++i]); break;
					case 'b': bedFile = new File(args[++i]); break;
					case 'p': peakShift = Integer.parseInt(args[++i]); break;
					case 'h': printDocs(); System.exit(0);
					default: Misc.printExit("\nProblem, unknown option! " + mat.group());
					}
				}
				catch (Exception e){
					Misc.printExit("\nSorry, something doesn't look right with this parameter: -"+test+"\n");
				}
			}
		}

		//look for point directories
		if (treatmentPointDirs == null || treatmentPointDirs[0].isDirectory() == false) Misc.printExit("\nError: cannot find your treatment PointData directories(s)!\n");
		//only one directory look deeper
		if (treatmentPointDirs.length == 1){
			File[] otherDirs = IO.extractOnlyDirectories(treatmentPointDirs[0]);
			if (otherDirs != null && otherDirs.length > 0) treatmentPointDirs = otherDirs;
		}

		//control data
		if (controlPointDirs != null){
			if (controlPointDirs[0].isDirectory() == false) Misc.printExit("\nError: cannot find your control PointData directories(s)!\n");
			//only one directory look deeper
			if (controlPointDirs.length == 1){
				File[] otherDirs = IO.extractOnlyDirectories(controlPointDirs[0]);
				if (otherDirs != null && otherDirs.length > 0) controlPointDirs = otherDirs;
			}
		}
		//set half peak shift
		halfPeakShift = (int)Math.round( ((double)peakShift)/2 );

		//need R
		if (fullPathToR == null || fullPathToR.exists() == false) Misc.printExit("\nPlease enter the full path to R, ie '-r /usr/bin/R'\n");

		//look for and or create the save directory
		if (saveDirectory == null) Misc.printExit("\nError: enter a directory name to save results.\n");
		if (saveDirectory.exists() == false) saveDirectory.mkdir();
		
		//look for bed file
		if (refSeqFile == null && bedFile == null){
			Misc.printExit("\nPlease enter a regions file to use in scoring regions.\n");
		}

	}	

	public static void printDocs(){
		System.out.println("\n" +
				"**************************************************************************************\n" +
				"**                         Defined Region Scan Seqs:  Jan 2009                      **\n" +
				"**************************************************************************************\n" +
				"Takes unshifted chromosome specific PointData xxx.bar.zip files and extracts scores\n" +
				"under each region to calculate several statistics including a binomial p-value and\n" +
				"Storey q-value FDR.  If a gene table is provided, scores under each exon are summed to\n" +
				"give a whole gene summary. Control data is optional but recommended.\n\n" +

				"Options:\n"+
				"-s Save directory, full path.\n"+
				"-t Treatment PointData directories, full path, comma delimited. These should\n" +
				"       contain unshifted strand and chromosome specific xxx+/-.bar.zip files. \n" +
				"-c Control PointData directories, ditto. \n" +
				"-p Peak shift, average distance between + and - strand peaks. Will be used to shift\n" +
				"       the PointData by 1/2 the peak shift, defaults to 0. See the PeakShiftFinder app. \n"+
				"-r Full path to R loaded with Storey's q-value library, defaults to '/usr/bin/R',\n" +
				"       see http://genomics.princeton.edu/storeylab/qvalue/\n"+
				"-u UCSC RefFlat or RefSeq gene table file, full path. See,\n"+
				"       http://genome.ucsc.edu/cgi-bin/hgTables\n"+
				"-b (Or) a bed file (chr, start, stop,...), full path, See,\n" +
				"       http://genome.ucsc.edu/FAQ/FAQformat#format1\n"+

				"\n"+

				"Example: java -Xmx1500M -jar pathTo/USeq/Apps/DefinedRegionScanSeqs -t\n" +
				"      /Data/PolIIRep1/,/Data/PolIIRep2/ -c /Data/Input1/,Data/Input2/ -b\n" +
				"      /Anno/miRNAs.bed -s /Results/NovelMiRNAs/ \n\n" +

		"**************************************************************************************\n");

	}
}
